package com.zhangcuishan.blog.controller;

import com.alibaba.fastjson.JSONObject;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.zhangcuishan.blog.util.EventStreamRestTemplate;
import lombok.extern.slf4j.Slf4j;
import org.apache.tomcat.util.http.fileupload.IOUtils;
import org.springframework.http.*;
import org.springframework.http.client.ClientHttpRequest;
import org.springframework.stereotype.Controller;
import org.springframework.ui.Model;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.CrossOrigin;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestMethod;
import org.springframework.web.client.ResponseExtractor;
import org.springframework.web.client.RestTemplate;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletResponse;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.UncheckedIOException;
import java.util.Locale;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

/**
 * @author wuq
 * @Time 2022-12-19 15:02
 * @Description
 */
@CrossOrigin(origins = "*", maxAge = 3600)
@Controller
@Slf4j
public class SSEController {

    @Resource
    private RestTemplate restTemplate;

    @RequestMapping(value = "/sse", method = RequestMethod.GET)
    public String home(Locale locale, Model model) {
        return "sse";
    }

    @RequestMapping(value = "/push")
    public void push(HttpServletResponse httpServletResponse) throws IOException {
        httpServletResponse.setContentType("text/event-stream");
        httpServletResponse.setCharacterEncoding("utf-8");

        String url = "http://39.105.215.32/chatmt-api/v1/llmchat";
        EventStreamRestTemplate restTemplate = new EventStreamRestTemplate();
        ResponseExtractor<String> responseExtractor = response -> {
            try (BufferedReader reader = new BufferedReader(new InputStreamReader(response.getBody()))) {
                return reader.lines().collect(Collectors.joining(System.lineSeparator()));
            } catch (IOException e) {
                throw new UncheckedIOException(e);
            }
        };

        HttpHeaders headers = new HttpHeaders();
        MediaType type = MediaType.parseMediaType("application/json; charset=UTF-8");
        headers.setContentType(type);
        headers.add("Accept", MediaType.APPLICATION_JSON.toString());

        JSONObject jsonObj = new JSONObject();
        jsonObj.put("inputstr", "美国现任总统是谁？");

        HttpEntity<String> formEntity = new HttpEntity<String>(jsonObj.toString(), headers);

        String result = restTemplate.execute(url, HttpMethod.POST, null, responseExtractor);
        System.out.println(result);


//        PrintWriter writer = response.getWriter();
//        BufferedReader bufferedReader;
//        try {
//            bufferedReader = new BufferedReader(new InputStreamReader(responseEntity.getBody().getInputStream()));
//            String line;
//            while ((line = bufferedReader.readLine()) != null && !(ChatGpt3dot5Request.STREAM_MESSAGE_PREFIX + "[DONE]").equals(line)) {
//                String message = getMessageFromLine(line, ChatGpt3dot5Request.STREAM_MESSAGE_PREFIX);
//                System.out.println(message);
////                writer.write(message);
////                writer.flush();
//            }
//        } catch (IOException e) {
//            throw new RuntimeException(e);
//        }


        for (int i = 0; i < 5; i++) {
            // 指定事件标识  event: 这个为固定格式
            httpServletResponse.getWriter().write("event:me\n");
            // 格式：data: + 数据 + 2个回车
            httpServletResponse.getWriter().write("data:" + i + "\n\n");
            httpServletResponse.getWriter().flush();

            try {
                TimeUnit.SECONDS.sleep(1);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
    }

    private static String getMessageFromLine(String line, String prefix) throws JsonProcessingException {
        if (!StringUtils.hasLength(line.trim())) {
            return "";
        }
        String messageJsonStr = line.substring(line.indexOf(prefix) + prefix.length());
        ObjectMapper objectMapper = new ObjectMapper();
        ObjectNode objectNode = (ObjectNode) objectMapper.readTree(messageJsonStr);
        JsonNode messageNode = objectNode.get("choices").get(0).get("delta").get("content");
        if (messageNode != null) {
            return messageNode.asText();
        } else {
            return "";
        }
    }

    @RequestMapping(value = "/yourEndPoint")
    public void processRequest(HttpServletResponse response) {
        RestTemplate restTemplate = new RestTemplate();

        String fileToDownloadUri= "http://39.105.215.32/chatmt-api/v1/llmchat";

        response.setStatus(HttpStatus.OK.value());
        restTemplate.execute(
                fileToDownloadUri,
                HttpMethod.POST,
                (ClientHttpRequest requestCallback) -> {},
                responseExtractor -> {
                    IOUtils.copy(responseExtractor.getBody(), response.getOutputStream());
                    return null;
                });
    }
}
